import os
import Utility.Stl as Stl
import Utility.tetrahedronTXT as Txt
from DEAD.AutoDecoder.Generator.GDS import GDS, save_as_vtk_from_npz
import torch
import numpy as np

'''
def calculate_mdf_from_cloud_points(bs_points_cloud, region_points_cloud):
    # 计算MDF值(无符号最小距离场)
    mdf = np.array([])
    with torch.no_grad():
        tensor_bs_points_cloud = torch.from_numpy(bs_points_cloud).cuda()
        num_rows_bs= tensor_bs_points_cloud.size(0)
        tensor_region_points_cloud = torch.from_numpy(region_points_cloud).cuda()
        num_rows_region = tensor_region_points_cloud.size(0)
        for i in range(num_rows_region):
            tensor_region_points_cloud_i = tensor_region_points_cloud[i].repeat(num_rows_bs, 1)
            mdf_i=torch.min(torch.sqrt(torch.sum((tensor_bs_points_cloud-tensor_region_points_cloud_i) ** 2, dim=1)))
            mdf = np.append(mdf, mdf_i.cpu().numpy())
    mdf = mdf.reshape(num_rows_region, 1)
    return mdf
'''

def calculate_mdf_from_cloud_points(bs_points_cloud, region_points_cloud):
    # 将数据转移到GPU
    tensor_bs_points_cloud = torch.from_numpy(bs_points_cloud).cuda()
    tensor_region_points_cloud = torch.from_numpy(region_points_cloud).cuda()

    # 使用cdist计算所有点对之间的距离矩阵
    distances = torch.cdist(tensor_region_points_cloud, tensor_bs_points_cloud)

    # 计算每个区域点的最小距离
    mdf = torch.min(distances, dim=1).values

    # 将结果转移到CPU并转换为numpy数组
    mdf = mdf.cpu().numpy().reshape(-1, 1)

    return mdf

def generate_grain_from_stl_and_txt(grain_index : int, data_read_path : str, data_write_path : str, points_num=5000, need_vtk : bool = False):
    # 生成编号为grain_index的药型数据（STL -> npz）
    # 在data_path路径下查找STL格式的药型数据，然后生成npz格式的药型数据（npz格式由GDS.py定义）

    # 确保data_path路径存在
    os.makedirs(data_write_path, exist_ok=True)

    # 提取信息stl文件名词信息（例如0#N6L3215M20BS）
    extraction = Stl.extract_info_from_stl_filename(os.listdir(data_read_path), grain_index)
    # 读取STL文件，获取初始燃面点云
    bs_points_cloud = Stl.read_stl_by_given_point_num(f"{data_read_path}/{extraction['stl_full_name']}",total_points=points_num)

    # 读取STL文件，获取初始燃面点云(加密版本)，用于计算MDF
    #bs_points_cloud_fine = Stl.read_stl_by_given_point_num(f"{data_read_path}/{extraction['stl_full_name']}",total_points=4*points_num)

    # 读取TXT文件，获取空腔区域点云
    cr_points_cloud = Txt.read_txt_by_given_point_num(f"{data_read_path}/{extraction['cr_txt_full_name']}",total_points=points_num)

    # 读取TXT文件，获取固体区域点云
    sr_points_cloud = Txt.read_txt_by_given_point_num(f"{data_read_path}/{extraction['sr_txt_full_name']}",total_points=points_num)

    # 创建GDS对象
    gds = GDS(lx=1.0,
            ly=1.0,
            lz=extraction["lz"],
            nx_points=1,
            ny_points=1,
            nz_points=1,
            n_slots=extraction["n_slots"],
            m =extraction["m"])
    
    # 设置燃面坐标
    gds.burning_surface_coordinate = bs_points_cloud

    # 设置空腔区域点云坐标
    gds.cavity_region_coordinate = cr_points_cloud

    # 设置固体区域点云坐标
    gds.solid_region_coordinate = sr_points_cloud

    # 计算空腔区域的SDF值(有符号最小距离场，都是负数)
    gds.cavity_region_sdf = -calculate_mdf_from_cloud_points(bs_points_cloud, cr_points_cloud)

    # 计算固体区域的SDF值(有符号最小距离场，都是正数)
    gds.solid_region_sdf = calculate_mdf_from_cloud_points(bs_points_cloud, sr_points_cloud)

    # 保存为npz格式
    gds.save_as_npz(f"{data_write_path}/{grain_index}")

    if need_vtk:
        # 转换为VTK格式（仅用于可视化）
        save_as_vtk_from_npz(f"{data_write_path}/{grain_index}",
                            f"{data_write_path}/{grain_index}")
        
    print(f"Finish Generating Grain: {grain_index}")


if __name__ == '__main__':
    #print(os.listdir("DEAD/AutoDecoder/Grains/G9"))
    #Stl.extract_info_from_stl_filename(os.listdir("DEAD/AutoDecoder/Grains/G9"), 0)
    generate_grain_from_stl_and_txt(0, "DEAD/AutoDecoder/Grains/G9", "DEAD/AutoDecoder/Grains/X", 5000, False)